import numpy as np


# 数据加载
def read_local_MNIST_dataset(filename=r".\mnist.npz"):
    f = np.load(filename)
    x_train, y_train = f['x_train'], f['y_train']
    x_test, y_test = f['x_test'], f['y_test']
    f.close()
    return (x_train, y_train), (x_test, y_test)


# n*m格子分割图片，得到n*m个手写特征，
def split_pic(image_set, n, m, threshold):
    chara_imgs = []
    for img in image_set:
        chara_vector = np.zeros((n, m))
        for i in range(n):
            for j in range(m):
                chara_vector[i, j] += np.sum(
                    img[i * (28 // n):(i + 1) * (28 // n), j * (28 // n):(j + 1) * (28 // m)] > threshold)
        chara_imgs.append(chara_vector.reshape(chara_vector.shape[0] * chara_vector.shape[1]))
    result = np.array(chara_imgs)
    # print('image_matrix shape -->', result.shape)
    return result  # 返回一个二维矩阵，每行为一张图片的特征


# 欧氏距离计算
# 对两个n维特征值进行计算并返回结果
class o_distance_classifier():
    def __init__(self, dim, class_num):
        self.class_num = class_num
        self.dim = dim
        self.class_center = np.zeros([dim, class_num])
        self.sample_count = [0] * class_num

    def train(self, train_images, train_labels):
        for img, y in zip(train_images, train_labels):
            self.class_center[:, y] += img
            self.sample_count[y] += 1
        self.class_center = self.class_center / self.sample_count

    def predict(self, x):
        x = x.T
        distance = [0] * self.class_num
        for i in range(self.class_num):
            distance[i] = np.dot((x - self.class_center[:, i]).T, (x - self.class_center[:, i]))  # 计算到各类中心的距离
        return np.argmin(distance)

    def save_param(self, filepath):
        param = {}
        param['class_center'] = self.class_center
        np.save(filepath, param)

    def load_param(self, filepath):
        param = np.load(filepath).item()
        self.class_center = param['class_center']


# 马氏距离计算
# 对两个n维特征值进行计算并返回结果
class m_distance_classifier():
    def __init__(self, dim, class_num):
        self.class_num = class_num
        self.dim = dim
        self.class_center = np.zeros([dim, class_num])
        self.sample_count = [0] * class_num
        self.cov_matrix = None
        self.cov_matrix_inverse = None

    def train(self, train_images, train_labels):
        for img, y in zip(train_images, train_labels):
            self.class_center[:, y] += img
            self.sample_count[y] += 1
        self.class_center = self.class_center / self.sample_count
        self.cov_matrix = np.cov(train_images, rowvar=0)
        self.cov_matrix_inverse = np.linalg.inv(self.cov_matrix)
        print("self.cov_matrix.shape-->", self.cov_matrix.shape)

    def predict(self, x):
        x = x.T
        distance = [0] * self.class_num
        for i in range(self.class_num):
            temp = np.dot((x - self.class_center[:, i]).T, self.cov_matrix_inverse)
            distance[i] = np.dot(temp, (x - self.class_center[:, i]))  # 计算到各类中心的距离
        return np.argmin(distance)

    def save_param(self, filepath):
        param = {}
        param['class_center'] = self.class_center
        param['cov_matrix_inverse'] = self.cov_matrix_inverse
        np.save(filepath, param)

    def load_param(self, filepath):
        param = np.load(filepath).item()
        self.class_center = param['class_center']
        self.cov_matrix_inverse = param['cov_matrix_inverse']


# PCA实现
class PCA_transformer():
    def __init__(self, chara_X):
        chara_X = chara_X.T
        self.R_matrix = np.dot(chara_X, chara_X.T)
        self.C_matrix = None
        self.eigenvalue = None
        self.m = None

    def generate(self, con_rate):
        e, v = np.linalg.eig(self.R_matrix)  # 计算特征值和特征向量
        sort_args = np.argsort(e)[::-1]  # 从大到小排序
        self.C_matrix = v[:, sort_args]
        self.eigenvalue = e[sort_args]
        # 根据贡献率计算m
        contr = np.sum(self.eigenvalue) * con_rate
        sum = 0
        for m, value in enumerate(self.eigenvalue):
            sum += value
            if sum > contr:
                break
        self.m = m
        con_rate = np.sum(self.eigenvalue[:self.m + 1]) / np.sum(self.eigenvalue)
        print("选择前%d个特征，贡献率：%.2f%%" % ((self.m + 1), con_rate * 100))
        print("self.eigenvalue.shape-->", self.eigenvalue.shape)
        # print("self.eigenvalue-->\n", self.eigenvalue)
        print("self.C_matrix.shape-->", self.C_matrix.shape)
        # print("self.C_matrix-->\n", self.C_matrix)

    def tran(self, x):
        # 功能：将特征向量x转化为y，y的前m个分量的贡献率不低于con_rate
        x = x.T
        result = np.dot(self.C_matrix[:self.m + 1, :], x)  # 只取Y的前m个分量代入计算
        return result.T

    def save_para(self, path_name):
        pass


def o_classification(x_train, y_train, x_test, y_test):
    ##　欧式距离分类
    o_classifier = o_distance_classifier(dim=x_train.shape[1], class_num=10)
    # 训练
    o_classifier.train(x_train, y_train)
    # 测试
    right_num = 0
    for img, label in zip(x_test, y_test):
        predict = o_classifier.predict(img)
        # print(predict, '======', label)
        if predict == label:
            right_num += 1
    print("right_num-->", right_num)
    print("len(y_test)-->", len(y_test))
    right_rate = right_num / len(y_test)
    print("欧式距离分类的准确率:%.2f%%" % (right_rate * 100))


def m_classification(x_train, y_train, x_test, y_test):
    ##　马式距离分类
    m_classifier = m_distance_classifier(dim=x_train.shape[1], class_num=10)
    # 训练
    m_classifier.train(x_train, y_train)
    # 测试
    right_num = 0
    for img, label in zip(x_test, y_test):
        predict = m_classifier.predict(img)
        if predict == label:
            right_num += 1
    print("right_num-->", right_num)
    print("len(y_train)-->", len(y_test))
    right_rate = right_num / len(y_test)
    print("马式距离分类的准确率:%.2f%%" % (right_rate * 100))


if __name__ == '__main__':
    ## 加载数据
    (x_train, y_train), (x_test, y_test) = read_local_MNIST_dataset()
    print("图片：")
    for i in range(28):
        for j in range(28):
            print(x_train[0][i][j], end='\t')
        print('')

    ##分割图片,减少特征维数
    sx_train = split_pic(x_train, 4, 4, 128)
    sx_test = split_pic(x_test, 4, 4, 128)
    print("sx_train.shape-->", sx_train.shape)
    print("y_train.shape-->", y_train.shape)
    print("sx_train[0]-->", sx_train[0])

    print("=" * 30, '使用16维特征进行分类', "=" * 30)
    ##　欧式距离分类
    o_classification(sx_train,y_train,sx_test,y_test)
    ##　马式距离分类
    m_classification(sx_train,y_train,sx_test,y_test)

    ## PCA主成分分析
    print("=" * 30, "PCA主成分分析", "=" * 30)
    pca_transformer = PCA_transformer(sx_train)
    pca_transformer.generate(con_rate=0.95)
    transx_train = pca_transformer.tran(sx_train)
    transx_test = pca_transformer.tran(sx_test)
    print("transx_train.shape-->", transx_train.shape)
    print("transx_test.shape-->", transx_test.shape)
    print("transx_train[0]-->", transx_train[0])

    o_classification(transx_train, y_train, transx_test, y_test)
    m_classification(transx_train, y_train, transx_test, y_test)


